/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: hdds op remote look up execute.
 * Author: liutianming
 * Create: 2022-11-28
 */

#ifndef OP_HDDS_REMOTE_LOOKUP_EXECUTOR_H
#define OP_HDDS_REMOTE_LOOKUP_EXECUTOR_H

#include <map>
#include <unordered_map>
#include <atomic>
#include <memory>
#include "hdds_param.h"
#include "op_embedding_service_executor.h"
#include "hccl_operator.h"
#include "transport_base_pub.h"
#include "aicpu_sharder.h" // aicpu sharding
#include "auto_buffer.h"
#include "con_unique_mappers_mgr.h"

namespace hccl {

using KeyPosTable = struct KeyPosTableDef {
    u64 psOffset;
    u32 pos;
    KeyPosTableDef() : psOffset(0), pos(0) {}
};

using KeyPsCount = struct KeyPsCountDef {
    u64 psOffset;
    u32 keyCount;
    KeyPsCountDef() : psOffset(0), keyCount(0) {}
};

struct RecoverValuePara {
    void *value = nullptr;
    u64 valueItemSize = 0;
    void *valueTransferMem = nullptr;
    bool isPipelineExec = false;
};

struct PrepareSendKeysBufPara {
    s64 *keys = nullptr;    // input
    u32 keyMaxNum = 0;      // input
    s64 *keyTransferMem = nullptr; // output
};

class HdcsOpRemoteLookUpExecutor : public OpEmbeddingServiceExecutor {
public:
    static constexpr u32 DOUBLE = 2;

    explicit HdcsOpRemoteLookUpExecutor();
    ~HdcsOpRemoteLookUpExecutor() override;
    HcclResult ExecuteOp(HcclOpPtr opPtr) override;
    HcclResult Init(std::vector<u32> &psRankId, u64 pipelineKeyNum = 0);
    HcclResult DeInit();

private:

    template <typename T, typename TM>
    using ConUniqueMappersInfo = typename ConUniqueMappersMgr<T, TM>::ConUniqueMappersInfo;

    HcclResult ExecuteTask(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult TaskDropDuplicates(OpTaskPtr &task);
    HcclResult TaskKeysDuplicatesShard(OpTaskPtr &task);
    HcclResult TaskSendKeys(OpTaskPtr &task);
    HcclResult TaskWaitLookupFinish(OpTaskPtr &task);
    HcclResult TaskRecoverValueShard(OpTaskPtr &task);

    HcclResult WaitSendKeyFinish(OpTaskPtr &task);
    HcclResult Send(PartitionMapTask *params, u32 rankId, void *buffer, u64 count, HcclDataType dataType,
        std::vector<HcclRequest> &requestArray) const;
    HcclResult WaitSomeAll(std::vector<HcclRequest> &requestArray);
    HcclResult RecvRespones(const PartitionMapTask *params, std::pair<u32, u32> &rankIdPsIdPair,
        std::vector<HcclRequest> &requestArray, std::unordered_map<u32, s32> &recvCounts);
    HcclResult TaskRestUnique(OpTaskPtr &task);

    HcclResult KeyReduce(OpTaskPtr &task);
    HcclResult KeyReduceWithUniqued(PartitionMapTask &partMapTask);
    HcclResult FillCounterWhileUniqued(PartitionMapTask &partMapTask);
    void PrepareMappingInfo(PartitionMapTask &partMapTask);
    HcclResult GetPsIdInfo(PartitionMapTask &partMapTask);
    HcclResult ReduceKeyAndGetMappingMatrix(PartitionMapTask &params);
    HcclResult ReduceKeyAndGetMappingMatrixWithNoUnique(PartitionMapTask &partMapTask);
    HcclResult FetchUniquedCounts();
    HcclResult AicpuRecoverValue(OpTaskPtr &task, std::function<bool(const u32&)> filtered);
    HcclResult ParralAicpuRecoverValue(OpTaskPtr &task);
    HcclResult ParralGetGetBatchPsIds(PartitionMapTask &partMapTask);

    u32 CalThreadNumEachPs(u32 keyMaxNum);
    void PrintLookupRecvSummary(PartitionMapTask *params, u32 psId);

    std::unordered_map<u32, Transport *> rankTransportMap_;
    HcclResult PrepareEnvelopeForRoce(PartitionMapTask *&params, void *buffer, u64 count,
        HcclDataType dataType, u32 psId);

    std::map<OpTaskType, std::function<HcclResult(OpTaskPtr &)>> taskExeFuncTable_;
    u32 keyMaxNum_ = 0; // 维护当前实例psIdsMem_内存的最大值
    u64 pipelineKeyNum_ = 0;
    u32 psSize_ = 0;
    std::vector<u32> psRankId_;
    std::unordered_map<u32, u32> rankId2PsId_;

    // 记录多线程去重后的{threadId, psId, 在sendbuf上的偏移位置, keysize}
    std::vector<std::vector<KeyPsCount>> keyPsSizeAfterUnique_; // index:threadId index:psId value:start offset
    std::vector<u32> keyOffsetForShard_; // 通过给每个thread分配的key num计算出对应的偏移
    std::vector<std::vector<u32>> psKeyNumForShard_;
    std::vector<std::vector<u64>> psOffsetForShard_;
    std::vector<std::vector<u32>> threadkeyNumForShard_; // 每个线程要处理的key num以及对应的key index
    u32 *psIdsMem_ = nullptr;

    // sort key by psId之后，记录psId所对应的内存相对于keyTransferMem的偏移值
    std::vector<u64> psOffset_;
    // 记录要发给每个ps实际的key数量，初始值为0，在去重过程累加以维护恢复矩阵
    std::vector<u32> psKeyNum_;

    std::vector<KeyPosTable> keyMapPsOffset_; // <keys pos， <psoffset, pos>> 恢复矩阵

    u32 shardExecMaxCoreNum_ = 1;
    std::vector<RecoverValuePara> recoverValueParas_;
    std::vector<PrepareSendKeysBufPara> prepareSendKeysBufParas_;
    std::vector<HcclRequest> requestArray_; // send keys的request放在流水最后去做

    // 去重的数据类型为s64，映射索引Matrix数据类型为s32
    ConUniqueMappersMgr<s64, s32> conUniqueMappersMgr_{};
    AutoBuffer<PsBufferInfo> psOffsetAndNum_{};
    AutoBuffer<s32> mappingMatrix_{};
};

}

#endif
